from elasticsearch import AsyncElasticsearch
import asyncio

async def setup_optimized_index(es: AsyncElasticsearch, index_name: str):
    """创建优化的ES索引"""
    
    optimized_settings = {
        "settings": {
            "index": {
                "number_of_shards": 3,
                "number_of_replicas": 1,
                "refresh_interval": "30s",  # 降低刷新频率，提高性能
                "codec": "best_compression"
            },
            "analysis": {
                "analyzer": {
                    "text_analyzer": {
                        "type": "custom",
                        "tokenizer": "standard",
                        "filter": ["lowercase", "asciifolding"]
                    }
                }
            }
        },
        "mappings": {
            "properties": {
                "filename": {"type": "text", "analyzer": "text_analyzer"},
                "file_url": {"type": "keyword"},
                "suffix": {"type": "keyword"},
                "fileclass": {"type": "keyword"},
                "keywords": {
                    "type": "keyword",
                    "ignore_above": 256
                },
                # 优化向量字段配置，启用HNSW算法
                "title_vector": {
                    "type": "dense_vector",
                    "dims": 1024,
                    "index": True,
                    "similarity": "cosine",
                    "index_options": {
                        "type": "hnsw",
                        "m": 16,              # 每个节点的连接数，增加可提高准确度但降低速度
                        "ef_construction": 100 # 构建索引时考虑的邻居数
                    }
                },
                "summary_text_vector": {
                    "type": "dense_vector",
                    "dims": 1024,
                    "index": True,
                    "similarity": "cosine",
                    "index_options": {
                        "type": "hnsw",
                        "m": 16,
                        "ef_construction": 100
                    }
                },
                "query_vector_1": {
                    "type": "dense_vector",
                    "dims": 1024,
                    "index": True,
                    "similarity": "cosine",
                    "index_options": {
                        "type": "hnsw",
                        "m": 16,
                        "ef_construction": 100
                    }
                },
                "query_vector_2": {"type": "dense_vector", "dims": 1024, "index": True, "similarity": "cosine"},
                "query_vector_3": {"type": "dense_vector", "dims": 1024, "index": True, "similarity": "cosine"},
                "title": {"type": "text", "analyzer": "text_analyzer"},
                "summary_text": {"type": "text", "analyzer": "text_analyzer"},
                "query_1": {"type": "text", "analyzer": "text_analyzer"},
                "query_2": {"type": "text", "analyzer": "text_analyzer"},
                "query_3": {"type": "text", "analyzer": "text_analyzer"},
                "file_status": {"type": "keyword"},
                "company": {"type": "keyword"},
                # 新增组合字段，便于全文搜索
                "combined_text": {
                    "type": "text",
                    "analyzer": "text_analyzer",
                    "fields": {
                        "keyword": {"type": "keyword"}
                    }
                },
                "metadata": {
                    "type": "object",
                    "enabled": True
                }
            }
        }
    }
    
    # 检查索引是否存在
    exists = await es.indices.exists(index=index_name)
    
    if exists:
        print(f"更新索引 {index_name} 的设置")
        # 关闭索引以更新设置
        await es.indices.close(index=index_name)
        # 更新索引设置
        await es.indices.put_settings(index=index_name, body=optimized_settings["settings"])
        # 重新打开索引
        await es.indices.open(index=index_name)
    else:
        print(f"创建新索引 {index_name}")
        # 创建新索引
        await es.indices.create(index=index_name, body=optimized_settings)
    
    return True 